"s3tablescatalog/cm-namespace-20241222"."cm_namespace"."lineorder"

import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from pyspark.conf import SparkConf  
from awsglue.job import Job

from pyspark.sql.functions import *
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# sc = SparkContext()
sc = SparkContext.getOrCreate()
glueContext = GlueContext(sc)
spark = glueContext.spark_session

def build_spark_session_s3t(namespace, warehouse_arn):
    spark = SparkSession.builder \
        .config(f"spark.sql.catalog.{namespace}", "org.apache.iceberg.spark.SparkCatalog") \
        .config(f"spark.sql.catalog.{namespace}.warehouse", warehouse_arn) \
        .config(f"spark.sql.catalog.{namespace}.catalog-impl", "software.amazon.s3tables.iceberg.S3TablesCatalog") \
        .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") \
        .getOrCreate()
    return(spark)

TARGET_DATE = '1992-10-01'
TARGET_START_DATE = TARGET_DATE
TARGET_END_DATE = '1992-10-31'

# Create Spark Session for S3 Tables 
warehouse_arn = "arn:aws:s3tables:us-east-1:123456789012:bucket/cm-namespace-20241222"
spark_s3t = build_spark_session_s3t("s3tablesbucket", warehouse_arn)

# Dynamic Frame from catalog
partition_predicate = f"lo_orderdate between '{TARGET_START_DATE}' and '{TARGET_END_DATE}'"
lineorder_p = glueContext.create_dynamic_frame.from_catalog(
    database="ssb",
    table_name="lineorder_p",
    push_down_predicate = partition_predicate,
    transformation_ctx="lineorder_p",
)
# lineorder_p.printSchema()
df = lineorder_p.toDF()
# df.printSchema()
# df.show()
# print(df.count())

# # Create namespace in S3 Tables
# spark_s3t.sql(""" CREATE NAMESPACE IF NOT EXISTS s3tablesbucket.cm_namespace """)           

spark_s3t.sql("""
DROP TABLE IF EXISTS s3tablesbucket.cm_namespace.`lineorder` PURGE
  """)

spark_s3t.sql("""
CREATE TABLE IF NOT EXISTS s3tablesbucket.cm_namespace.`lineorder` (
    `lo_orderkey` int,
    `lo_linenumber` int,
    `lo_custkey` int,
    `lo_partkey` int,
    `lo_suppkey` int,
    `lo_orderpriority` string,
    `lo_shippriority` string,
    `lo_quantity` int,
    `lo_extendedprice` int,
    `lo_ordertotalprice` int,
    `lo_discount` int,
    `lo_revenue` int,
    `lo_supplycost` int,
    `lo_tax` int,
    `lo_commitdate` string,
    `lo_shipmode` string,
    `lo_orderdate` string
) USING iceberg
PARTITIONED BY (`lo_orderdate`)
  """)

# # Use Namespace
# spark_s3t.sql(""" USE s3tablesbucket.cm_namespace """)

# # Show Tables
# spark_s3t.sql(""" SHOW TABLES """).show()

# Insert records into Iceberg table in S3 Tables
df.createOrReplaceTempView("tmp_lineorder")
spark_s3t.sql(""" 
  INSERT INTO s3tablesbucket.cm_namespace.`lineorder` 
  SELECT * FROM tmp_lineorder
""")

import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from pyspark.conf import SparkConf  
from awsglue.job import Job

from pyspark.sql.functions import *
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# sc = SparkContext()
sc = SparkContext.getOrCreate()
glueContext = GlueContext(sc)
spark = glueContext.spark_session

def build_spark_session_s3i(catalog_name, warehouse_path):
    spark = SparkSession.builder \
        .config("spark.sql.warehouse.dir", warehouse_path) \
        .config(f"spark.sql.catalog.{catalog_name}", "org.apache.iceberg.spark.SparkCatalog") \
        .config(f"spark.sql.catalog.{catalog_name}.warehouse", warehouse_path) \
        .config(f"spark.sql.catalog.{catalog_name}.catalog-impl", "org.apache.iceberg.aws.glue.GlueCatalog") \
        .config(f"spark.sql.catalog.{catalog_name}.io-impl", "org.apache.iceberg.aws.s3.S3FileIO") \
        .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") \
        .getOrCreate()
    return(spark)

TARGET_DATE = '1992-10-01'
TARGET_START_DATE = TARGET_DATE
TARGET_END_DATE = '1992-10-31'

# Create Spark Session for Iceberg Tables on Glue
i_catalog_name = "glue_catalog"
i_bucket_name = "cm-datalaker-20241222"
i_bucket_prefix = "lineorder"
i_database_name = "ssb_iceberg"
i_table_name = "lineorder"
i_warehouse_path = f"s3://{i_bucket_name}/{i_bucket_prefix}"
spark_s3t = build_spark_session_s3i(i_catalog_name, i_warehouse_path)

# Dynamic Frame from catalog
partition_predicate = f"lo_orderdate between '{TARGET_START_DATE}' and '{TARGET_END_DATE}'"
lineorder_p = glueContext.create_dynamic_frame.from_catalog(
    database="ssb",
    table_name="lineorder_p",
    push_down_predicate = partition_predicate,
    transformation_ctx="lineorder_p",
)
# lineorder_p.printSchema()
df = lineorder_p.toDF()
# df.printSchema()
# df.show()
# print(df.count())

spark_s3t.sql(f"""
DROP TABLE IF EXISTS {i_catalog_name}.{i_database_name}.{i_table_name} PURGE
  """)

spark_s3t.sql(f"""
CREATE TABLE IF NOT EXISTS {i_catalog_name}.{i_database_name}.{i_table_name} (
    `lo_orderkey` int,
    `lo_linenumber` int,
    `lo_custkey` int,
    `lo_partkey` int,
    `lo_suppkey` int,
    `lo_orderpriority` string,
    `lo_shippriority` string,
    `lo_quantity` int,
    `lo_extendedprice` int,
    `lo_ordertotalprice` int,
    `lo_discount` int,
    `lo_revenue` int,
    `lo_supplycost` int,
    `lo_tax` int,
    `lo_commitdate` string,
    `lo_shipmode` string,
    `lo_orderdate` string
) USING iceberg
PARTITIONED BY (`lo_orderdate`)
  """)

# Insert records into Iceberg table in S3 Tables
df.createOrReplaceTempView("tmp_lineorder")
spark_s3t.sql(f""" 
  INSERT INTO {i_catalog_name}.{i_database_name}.{i_table_name} 
  SELECT * FROM tmp_lineorder
""")

select c_city, s_city, d_year, sum(lo_revenue) AS revenue
from "s3tablescatalog/cm-namespace-20241222"."cm_namespace"."lineorder"
inner join ssb.customer on lo_custkey = c_custkey
inner join ssb.supplier on lo_suppkey = s_suppkey
inner join ssb.dwdate on replace(lo_orderdate, '-', '') = d_datekey
GROUP BY c_city, s_city, d_year
ORDER BY d_year ASC, revenue DESC

select c_city, s_city, d_year, sum(lo_revenue) AS revenue
from ssb_iceberg.lineorder
inner join ssb.customer on lo_custkey = c_custkey
inner join ssb.supplier on lo_suppkey = s_suppkey
inner join ssb.dwdate on replace(lo_orderdate, '-', '') = d_datekey
GROUP BY c_city, s_city, d_year
ORDER BY d_year ASC, revenue DESC